from collections.abc import Iterable, Iterator
from typing import Any, Callable, Generic, TypeVar

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, DatasetWriter, pause, resume, skip
from zkl_aiutils_datasets.wrapping import wrap_dataset, wrap_dataset_iterator

SrcSample = TypeVar('SrcSample')
DstSample = TypeVar('DstSample')


class MappedDatasetWriter(DatasetWriter[SrcSample], Generic[SrcSample, DstSample]):
    def __init__(self, writer: DatasetWriter[DstSample], map_func: Callable[[SrcSample], DstSample]):
        self._writer = writer
        self._map_func = map_func

    @property
    def writer(self) -> DatasetWriter[DstSample]:
        return self._writer

    @property
    def map_func(self) -> Callable[[SrcSample], DstSample]:
        return self._map_func

    def write(self, sample: DstSample):
        sample = self._map_func(sample)
        self._writer.write(sample)

    def flush(self):
        self._writer.flush()

    def close(self):
        self._writer.close()


class MappedIterator(DatasetIterator[DstSample], Generic[DstSample, SrcSample]):
    def __init__(self, iterator: Iterator[DstSample] | Iterable[DstSample], map_func: Callable[[SrcSample], DstSample]):
        self._iterator = wrap_dataset_iterator(iterator)
        self._map_func = map_func

    def __next__(self) -> DstSample:
        src_sample = next(self._iterator)
        dst_sample = self._map_func(src_sample)
        return dst_sample

    def __skip__(self, samples_n: int):
        skip(self._iterator, samples_n)

    def __pause__(self) -> Any | None:
        return pause(self._iterator)


class MappedDataset(Dataset[DstSample], Generic[DstSample, SrcSample]):
    def __init__(self, dataset: Iterable[SrcSample], map_func: Callable[[SrcSample], DstSample]):
        self._dataset = wrap_dataset(dataset)
        self._map_func = map_func

    @property
    def dataset(self) -> Dataset[SrcSample]:
        return self._dataset

    @property
    def map_func(self) -> Callable[[SrcSample], DstSample]:
        return self._map_func

    def __len__(self) -> int:
        return len(self._dataset)

    def __getitem__(self, index: int, /) -> DstSample:
        return self._map_func(self._dataset[index])

    def __iter__(self) -> MappedIterator[DstSample, SrcSample]:
        return MappedIterator(iter(self._dataset), self._map_func)

    def __resume__(self, state: Any | None = None, /) -> MappedIterator[DstSample, SrcSample]:
        return MappedIterator(resume(self._dataset, state), self._map_func)
